from fastapi import APIRouter, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse
import time
import base64
import io
import asyncio
from typing import List, Literal, Union
from starlette.concurrency import run_in_threadpool

from ...dto.image import (
    ImageGenerationRequest,
    ImageGenerationResponse,
    GeneratedImageData,
    ImageGenPartialImage,
)
from ...handler.images import generate_image_entrypoint
from PIL import Image

# 创建路由器实例
router = APIRouter(prefix="/v1/images", tags=["images"])


def image_to_format(
    image: Image.Image, resp_format: Literal["b64_json", "url"]
) -> Union[str, GeneratedImageData]:
    """将PIL图像转换为指定格式的回复"""
    # 根据响应格式处理图像
    match resp_format:
        case "b64_json":
            # 转换为base64编码
            buffer = io.BytesIO()
            image.save(buffer, format="PNG")
            img_base64 = base64.b64encode(buffer.getvalue()).decode()
            return GeneratedImageData(b64_json=img_base64)
        case "url":
            # 模拟URL返回
            return GeneratedImageData(url="http://example.com/generated_image.png")
        case _:
            raise ValueError(f"Unsupported response format: {resp_format}")


async def stream_image_generation(image_futures: List, request: ImageGenerationRequest):
    """生成流式图像响应的异步生成器"""
    created_time = int(time.time())
    # https://platform.openai.com/docs/guides/image-generation?image-generation-model=gpt-image-1#generate-images
    # 需要参考一下, 功能对齐
    # 发送初始响应头
    initial_response = ImageGenPartialImage(
        b64_json="",
        background="auto",
        created_at=created_time,
        output_format="png",
        partial_image_index=0,  #
        quality=request.quality,
        size=request.size,
        type="image_generation.partial_image",
    )
    yield f"data: {initial_response.model_dump_json()}\n\n"
    # TODO: unimplemented, 需要实现图像生成的流式响应(包括partial的数据)
    yield "data: [DONE]\n\n"


@router.post(
    "/generations",
    response_model=ImageGenerationResponse,
    responses={
        200: {"model": ImageGenerationResponse},
    },
)
async def create_image_generation(
    request: ImageGenerationRequest, background_tasks: BackgroundTasks = None
) -> Union[ImageGenerationResponse, StreamingResponse]:
    """
    根据文本提示生成图像

    Args:
        request: 图像生成请求参数
        background_tasks: 后台任务（可选）

    Returns:
        ImageGenerationResponse 或 StreamingResponse: 生成的图像响应
    """
    # 验证请求参数
    # TODO: 中间件加上请求接入时的LOG
    if not request.prompt.strip():
        raise HTTPException(status_code=400, detail="Prompt cannot be empty")

    # 生成图像Future列表
    image_futures = generate_image_entrypoint(request)

    if request.stream:
        # 返回流式响应
        return StreamingResponse(
            stream_image_generation(image_futures, request),
            media_type="text/plain",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "Content-Type": "text/plain; charset=utf-8",
            },
        )
    else:
        # 非流式响应：等待所有图像生成完成
        images = await asyncio.gather(*image_futures)
        generated_images: List[GeneratedImageData] = []

        for i, image in enumerate(images):
            if isinstance(image, Exception):
                raise HTTPException(
                    status_code=500,
                    detail=f"Error during image {i} generation",
                )
            formatted_image = await run_in_threadpool(
                image_to_format, image, request.response_format
            )
            generated_images.append(formatted_image)

        # 构造响应
        response = ImageGenerationResponse(
            created=int(time.time()), data=generated_images
        )

        return response
